import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np

from my_models.base import Base

# 创建模拟的时间序列数据
def create_time_series_data(seq_length=10, num_samples=1000):
    time = np.linspace(0, 50, seq_length * num_samples)
    data = np.sin(time) + 0.05 * np.random.randn(len(time))
    X = []
    y = []
    for i in range(num_samples):
        X.append(data[i * seq_length:(i + 1) * seq_length])
        y.append(data[(i + 1) * seq_length])  # 预测下一个值
    return np.array(X), np.array(y)

class TimeSeriesDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.tensor(X, dtype=torch.float32).unsqueeze(-1)  # 添加通道维度
        self.y = torch.tensor(y, dtype=torch.float32)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

# 定义 LSTM 模型
class LSTMModel(nn.Module):
    def __init__(self, input_dim=1, hidden_dim=50, output_dim=1, num_layers=2):
        super(LSTMModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)
        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])  # 只取最后一个时间步的输出
        return out


seq_length = 10
num_samples = 1000
batch_size = 64
epochs = 100
learning_rate = 0.01

class LstmRegressor(Base):

    def __init__(self):
        # 设置参数

        self.X, self.y = create_time_series_data(seq_length, num_samples)
        self.model = LSTMModel()
        self.criterion = nn.MSELoss()
        self.optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)

    def train(self, x_train, y_train):
        dataset = TimeSeriesDataset(self.X, self.y)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

        for epoch in range(epochs):
            for batch_X, batch_y in dataloader:
                batch_X, batch_y = batch_X.to('cpu'), batch_y.to('cpu')  # 根据你的设备选择 'cuda' 或 'cpu'
                self.optimizer.zero_grad()
                outputs = self.model(batch_X)
                loss = self.criterion(outputs.squeeze(), batch_y)
                loss.backward()
                self.optimizer.step()
            if (epoch + 1) % 10 == 0:
                print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}')

        self.clf = self.regressor.fit(x_train, y_train)

    def valid(self, x_test, y_test):
        with torch.no_grad():
            test_inputs = torch.tensor(self.X[:10], dtype=torch.float32).unsqueeze(-1)  # 使用前10个样本测试
            predictions = self.model(test_inputs).squeeze().numpy()
            actuals = self.y[:10]
            for i in range(10):
                print(f"Predicted: {predictions[i]:.3f}, Actual: {actuals[i]:.3f}")

    def save_model(self, save_path):
        pass

    def load_model(self, model_path):
        # 测试模型
        pass





# 创建数据集
X, y = create_time_series_data(seq_length, num_samples)
dataset = TimeSeriesDataset(X, y)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# 初始化模型、损失函数和优化器


# 训练模型


